#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
继电器MQTT自动压力测试程序
用于持续测试ESP32继电器控制板的MQTT协议功能
协议格式：[0x55, 地址, 功能码, 数据1, 数据2, 数据3, 数据4, 校验和]
"""

import paho.mqtt.client as mqtt
import time
import sys
import struct
import threading
from datetime import datetime
import random

class RelayMqttStressTester:
    def __init__(self, broker_host="110.41.147.192", broker_port=1883, device_mac="3c8a1fd51e48"):
        """
        初始化MQTT压力测试客户端
        
        Args:
            broker_host: MQTT服务器地址
            broker_port: MQTT服务器端口
            device_mac: 设备MAC地址
        """
        self.broker_host = broker_host
        self.broker_port = broker_port
        self.device_mac = device_mac
        
        # MQTT主题
        self.publish_topic = f"/c/{device_mac}/get"  # 发送给设备
        self.subscribe_topic = f"/c/{device_mac}/send"  # 接收设备回复
        
        # 初始化MQTT客户端
        self.client = mqtt.Client()
        self.client.on_connect = self.on_connect
        self.client.on_message = self.on_message
        self.client.on_disconnect = self.on_disconnect
        
        # 测试统计
        self.test_stats = {
            'total_commands': 0,
            'successful_commands': 0,
            'failed_commands': 0,
            'timeout_commands': 0,
            'start_time': None,
            'last_response_time': None,
            'protocol_errors': 0,
            'response_times': [],  # 存储响应时间
            'test_types': {  # 各种测试类型的统计
                '读取状态': {'total': 0, 'success': 0, 'failed': 0, 'timeout': 0},
                '随机继电器开启': {'total': 0, 'success': 0, 'failed': 0, 'timeout': 0},
                '随机继电器关闭': {'total': 0, 'success': 0, 'failed': 0, 'timeout': 0},
                '所有继电器开启': {'total': 0, 'success': 0, 'failed': 0, 'timeout': 0},
                '所有继电器关闭': {'total': 0, 'success': 0, 'failed': 0, 'timeout': 0},
                '延时测试(开启)': {'total': 0, 'success': 0, 'failed': 0, 'timeout': 0},
                '延时测试(关闭)': {'total': 0, 'success': 0, 'failed': 0, 'timeout': 0},
            }
        }
        
        # 响应验证
        self.waiting_for_response = False
        self.response_received = False
        self.response_timeout = 15.0  # 15秒超时
        self.last_command_description = ""
        self.current_test_type = ""
        self.response_lock = threading.Lock()
        
        print(f"🔧 继电器MQTT压力测试器初始化")
        print(f"📡 MQTT服务器: {broker_host}:{broker_port}")
        print(f"🎯 设备MAC: {device_mac}")
        print(f"📤 发布主题: {self.publish_topic}")
        print(f"📥 订阅主题: {self.subscribe_topic}")
        print(f"⏱️  响应超时: {self.response_timeout}秒")
        print("-" * 80)

    def calculate_checksum(self, data):
        """计算校验和"""
        return sum(data) & 0xFF

    def create_protocol_packet(self, address=0x01, function_code=0x10, data1=0x00, data2=0x00, data3=0x00, data4=0x00):
        """
        创建8字节协议数据包
        
        Args:
            address: 设备地址
            function_code: 功能码
            data1-data4: 数据字节
            
        Returns:
            bytes: 8字节协议数据包
        """
        packet = [0x55, address, function_code, data1, data2, data3, data4, 0x00]
        packet[7] = self.calculate_checksum(packet[:7])
        return bytes(packet)

    def on_connect(self, client, userdata, flags, rc):
        """MQTT连接回调"""
        if rc == 0:
            print("✅ MQTT连接成功!")
            client.subscribe(self.subscribe_topic)
            print(f"📥 已订阅主题: {self.subscribe_topic}")
        else:
            print(f"❌ MQTT连接失败，错误码: {rc}")

    def on_message(self, client, userdata, msg):
        """MQTT消息接收回调"""
        with self.response_lock:
            if self.waiting_for_response:
                self.response_received = True
                self.test_stats['last_response_time'] = datetime.now()
                
                print(f"📨 收到设备回复 [{self.last_command_description}]:")
                print(f"   时间: {datetime.now().strftime('%H:%M:%S.%f')[:-3]}")
                print(f"   长度: {len(msg.payload)} 字节")
                
                # 解析8字节回复数据
                if len(msg.payload) == 8:
                    data = list(msg.payload)
                    print(f"   数据: {' '.join([f'0x{b:02X}' for b in data])}")
                    
                    # 验证协议格式
                    if data[0] == 0x22:  
                        function_code = data[2]
                        relay_state1 = data[5]  # 第16-9路继电器状态
                        relay_state2 = data[6]  # 第8-1路继电器状态
                        
                        print(f"   功能码: 0x{function_code:02X}")
                        print(f"   继电器状态: 高8位=0x{relay_state1:02X}({relay_state1:08b}) 低8位=0x{relay_state2:02X}({relay_state2:08b})")
                        print("   ✅ 响应格式正确")
                    else:
                        print(f"   ⚠️  协议头错误，期望0x55，实际0x{data[0]:02X}")
                        self.test_stats['protocol_errors'] += 1
                else:
                    print(f"   ⚠️  数据长度异常: {len(msg.payload)}字节")
                    print(f"   原始数据: {msg.payload.hex()}")
                    self.test_stats['protocol_errors'] += 1

    def on_disconnect(self, client, userdata, rc):
        """MQTT断开连接回调"""
        print(f"🔌 MQTT连接断开，错误码: {rc}")

    def connect(self):
        """连接MQTT服务器"""
        try:
            print(f"🔄 正在连接MQTT服务器...")
            self.client.connect(self.broker_host, self.broker_port, 60)
            self.client.loop_start()
            time.sleep(2)  # 等待连接建立
            return True
        except Exception as e:
            print(f"❌ 连接失败: {e}")
            return False

    def disconnect(self):
        """断开MQTT连接"""
        self.client.loop_stop()
        self.client.disconnect()
        print("👋 已断开MQTT连接")

    def send_command_with_verification(self, packet, description="", test_type=""):
        """发送命令并验证响应"""
        with self.response_lock:
            self.waiting_for_response = True
            self.response_received = False
            self.last_command_description = description
            self.current_test_type = test_type
        
        # 统计测试类型
        if test_type in self.test_stats['test_types']:
            self.test_stats['test_types'][test_type]['total'] += 1
        
        try:
            packet_hex = ' '.join([f'0x{b:02X}' for b in packet])
            print(f"\n📤 [{datetime.now().strftime('%H:%M:%S.%f')[:-3]}] 发送: {description}")
            print(f"   数据: {packet_hex}")
            
            # 统计发送命令
            self.test_stats['total_commands'] += 1
            
            result = self.client.publish(self.publish_topic, packet)
            if result.rc != mqtt.MQTT_ERR_SUCCESS:
                print(f"   ❌ 发送失败，错误码: {result.rc}")
                self.test_stats['failed_commands'] += 1
                if test_type in self.test_stats['test_types']:
                    self.test_stats['test_types'][test_type]['failed'] += 1
                with self.response_lock:
                    self.waiting_for_response = False
                return False
            
            # 等待响应
            start_time = time.time()
            while time.time() - start_time < self.response_timeout:
                with self.response_lock:
                    if self.response_received:
                        response_time = (time.time() - start_time) * 1000
                        print(f"   ✅ 命令执行成功，响应时间: {response_time:.1f}ms")
                        self.test_stats['successful_commands'] += 1
                        self.test_stats['response_times'].append(response_time)
                        if test_type in self.test_stats['test_types']:
                            self.test_stats['test_types'][test_type]['success'] += 1
                        self.waiting_for_response = False
                        return True
                time.sleep(0.1)
            
            # 超时
            print(f"   ⏰ 响应超时 ({self.response_timeout}秒)")
            self.test_stats['timeout_commands'] += 1
            if test_type in self.test_stats['test_types']:
                self.test_stats['test_types'][test_type]['timeout'] += 1
            with self.response_lock:
                self.waiting_for_response = False
            return False
            
        except Exception as e:
            print(f"   ❌ 发送异常: {e}")
            self.test_stats['failed_commands'] += 1
            if test_type in self.test_stats['test_types']:
                self.test_stats['test_types'][test_type]['failed'] += 1
            with self.response_lock:
                self.waiting_for_response = False
            return False

    def print_statistics(self):
        """打印详细的测试统计信息"""
        if self.test_stats['start_time']:
            elapsed = (datetime.now() - self.test_stats['start_time']).total_seconds()
            rate = self.test_stats['total_commands'] / elapsed if elapsed > 0 else 0
            success_rate = (self.test_stats['successful_commands'] / self.test_stats['total_commands'] * 100) if self.test_stats['total_commands'] > 0 else 0
            
            print(f"\n" + "=" * 80)
            print(f"📊 详细测试统计 (运行时间: {elapsed:.1f}秒)")
            print("=" * 80)
            
            # 总体统计
            print(f"🎯 总体统计:")
            print(f"   总命令数: {self.test_stats['total_commands']}")
            print(f"   成功: {self.test_stats['successful_commands']} ({success_rate:.1f}%)")
            print(f"   失败: {self.test_stats['failed_commands']}")
            print(f"   超时: {self.test_stats['timeout_commands']}")
            print(f"   协议错误: {self.test_stats['protocol_errors']}")
            print(f"   命令频率: {rate:.2f} 命令/秒")
            
            # 响应时间统计
            if self.test_stats['response_times']:
                response_times = self.test_stats['response_times']
                avg_time = sum(response_times) / len(response_times)
                min_time = min(response_times)
                max_time = max(response_times)
                print(f"   平均响应时间: {avg_time:.1f}ms")
                print(f"   最短响应时间: {min_time:.1f}ms")
                print(f"   最长响应时间: {max_time:.1f}ms")
            
            if self.test_stats['last_response_time']:
                print(f"   最后响应: {self.test_stats['last_response_time'].strftime('%H:%M:%S')}")
            
            # 各测试类型统计
            print(f"\n📋 各测试类型统计:")
            print("-" * 80)
            for test_type, stats in self.test_stats['test_types'].items():
                if stats['total'] > 0:
                    success_rate_type = (stats['success'] / stats['total'] * 100) if stats['total'] > 0 else 0
                    print(f"   {test_type}:")
                    print(f"     总计: {stats['total']} | 成功: {stats['success']} ({success_rate_type:.1f}%) | 失败: {stats['failed']} | 超时: {stats['timeout']}")
            
            print("=" * 80)

    # 测试函数集合 - 添加test_type参数
    def test_read_status(self):
        """测试读取继电器状态"""
        packet = self.create_protocol_packet(function_code=0x10)
        return self.send_command_with_verification(packet, "读取继电器状态", "读取状态")

    def test_open_relay(self, relay_num):
        """测试打开指定继电器"""
        if not (1 <= relay_num <= 14):
            print(f"❌ 继电器编号错误: {relay_num} (支持1-14)")
            return False
        packet = self.create_protocol_packet(function_code=0x12, data4=relay_num)
        return self.send_command_with_verification(packet, f"打开继电器{relay_num}", "随机继电器开启")

    def test_close_relay(self, relay_num):
        """测试关闭指定继电器"""
        if not (1 <= relay_num <= 14):
            print(f"❌ 继电器编号错误: {relay_num} (支持1-14)")
            return False
        packet = self.create_protocol_packet(function_code=0x11, data4=relay_num)
        return self.send_command_with_verification(packet, f"关闭继电器{relay_num}", "随机继电器关闭")

    def test_open_all(self):
        """测试打开所有继电器"""
        packet = self.create_protocol_packet(function_code=0x13, data4=0xFF)
        return self.send_command_with_verification(packet, "打开所有继电器", "所有继电器开启")

    def test_close_all(self):
        """测试关闭所有继电器"""
        packet = self.create_protocol_packet(function_code=0x13, data4=0x00)
        return self.send_command_with_verification(packet, "关闭所有继电器", "所有继电器关闭")

    def test_open_delay(self, relay_num, delay_ms):
        """测试打开延时关闭"""
        if not (1 <= relay_num <= 14):
            return False
        
        data1 = (delay_ms >> 16) & 0xFF
        data2 = (delay_ms >> 8) & 0xFF
        data3 = delay_ms & 0xFF
        
        packet = self.create_protocol_packet(function_code=0x21, data1=data1, data2=data2, data3=data3, data4=relay_num)
        return self.send_command_with_verification(packet, f"继电器{relay_num}开启{delay_ms}ms后关闭", "延时测试(开启)")

    def test_close_delay(self, relay_num, delay_ms):
        """测试关闭延时打开"""
        if not (1 <= relay_num <= 14):
            return False
        
        data1 = (delay_ms >> 16) & 0xFF
        data2 = (delay_ms >> 8) & 0xFF
        data3 = delay_ms & 0xFF
        
        packet = self.create_protocol_packet(function_code=0x22, data1=data1, data2=data2, data3=data3, data4=relay_num)
        return self.send_command_with_verification(packet, f"继电器{relay_num}关闭{delay_ms}ms后开启", "延时测试(关闭)")

def run_stress_test(tester):
    """运行压力测试"""
    print(f"\n🚀 开始压力测试... (按Ctrl+C停止)")
    print("=" * 80)
    
    tester.test_stats['start_time'] = datetime.now()
    
    test_cases = [
        # 基础功能测试
        ("读取状态", lambda: tester.test_read_status()),
        
        # 单个继电器测试
        ("随机继电器开启", lambda: tester.test_open_relay(random.randint(1, 14))),
        ("随机继电器关闭", lambda: tester.test_close_relay(random.randint(1, 14))),
        
        # 全部继电器测试
        ("所有继电器开启", lambda: tester.test_open_all()),
        ("所有继电器关闭", lambda: tester.test_close_all()),
        
        # 延时测试
        ("延时测试(开启)", lambda: tester.test_open_delay(random.randint(1, 14), random.randint(1000, 5000))),
        ("延时测试(关闭)", lambda: tester.test_close_delay(random.randint(1, 14), random.randint(1000, 5000))),
    ]
    
    cycle_count = 0
    
    try:
        while True:
            cycle_count += 1
            print(f"\n🔄 压力测试周期 #{cycle_count}")
            print("-" * 60)
            
            # 随机选择测试用例或按顺序执行
            if random.random() < 0.7:  # 70%概率随机测试
                test_name, test_func = random.choice(test_cases)
                test_func()
            else:  # 30%概率顺序测试
                for test_name, test_func in test_cases:
                    test_func()
                    time.sleep(0.5)  # 测试间隔
            
            # 每10个周期打印详细统计信息
            if cycle_count % 10 == 0:
                tester.print_statistics()
            
            # 测试间隔
            time.sleep(random.uniform(0.5, 2.0))
            
    except KeyboardInterrupt:
        print(f"\n\n⛔ 用户停止压力测试")
    except Exception as e:
        print(f"\n❌ 压力测试异常: {e}")
    
    # 最终统计
    print("\n" + "=" * 80)
    print("🏁 压力测试结束")
    tester.print_statistics()

def main():
    """主函数"""
    print("🎯 继电器MQTT自动压力测试程序")
    print("📝 请输入设备参数 (直接按Enter使用默认值)")
    
    # 获取用户输入
    broker = input("MQTT服务器地址 [110.41.147.192]: ").strip() or "110.41.147.192"
    port_str = input("MQTT端口 [1883]: ").strip() or "1883"
    mac = input("设备MAC地址 [3c8a1fd51e48]: ").strip() or "3c8a1fd51e48"
    
    try:
        port = int(port_str)
    except ValueError:
        print("❌ 端口号无效，使用默认值1883")
        port = 1883
    
    # 创建压力测试器
    tester = RelayMqttStressTester(broker, port, mac)
    
    # 连接MQTT服务器
    if not tester.connect():
        print("❌ 无法连接MQTT服务器，程序退出")
        return
    
    try:
        # 等待MQTT连接稳定
        print("⏳ 等待MQTT连接稳定...")
        time.sleep(3)
        
        # 运行压力测试
        run_stress_test(tester)
        
    except Exception as e:
        print(f"\n❌ 程序异常: {e}")
    finally:
        tester.disconnect()
        print("👋 程序结束")

if __name__ == "__main__":
    main() 